#include "ip_reassembly.h"
#include "IPv4Layer.h"
#include "IPv6Layer.h"
#include "PacketUtils.h"
#include <string.h>
#include "endian_portable.h"
#include "Packet.h"
#include "log/log.h"

namespace secsmart_tcpip
{

IPReassembly::IPFragData::IPFragData(uint8_t *data, size_t len,  uint16_t offset, bool last)
{
    this->offset = offset;
    this->last = last;
    this->data = new uint8_t[len];
    if (this->data != NULL) {
        memcpy(this->data, data, len);
    }
    this->data_len = len;
}

IPReassembly::IPFragData::~IPFragData() 
{
    if (data != NULL) {
        delete [] data;
    }
}

IPReassembly::IPFragStream::IPFragStream(IPFragStreamId id)
{
    this->id = id;
    expect_offset = 0;
    rcv_size = 0;
    total_size = 0;
    rcv_last = 0;
    rcv_first = 0;
    update_time = getSystemRunTime();
}
IPReassembly::IPFragStream::~IPFragStream() 
{
    while (!frag_list.empty()){
        IPFragData* frag = frag_list.front();
        delete frag;
        frag_list.pop_front();
    }
    frag_list.clear();
}

IPReassembly::IPReassembly(size_t streams_num, size_t frags_num, uint32_t timeout)
{
    m_stat.clear();

    m_cfg.detect_ooo_enable = true;
    m_cfg.max_frag_cache_num = (uint32_t)streams_num;
    m_cfg.max_stream_num = (uint32_t)frags_num;
    m_cfg.timeout = timeout; /*毫秒*/
    m_aging_count_once = 100; 
    m_last_aging_time = 0;
}

IPReassembly::~IPReassembly()
{
    LOG_TRACE("~IPReassembly List count:{} map count:{}", m_AgingList.size(), m_IPFragStreamMap.size());
    // 删除所有的IP分片链
    while(!m_AgingList.empty()) {
        // 取链表头
        auto listIter = m_AgingList.begin();
        deleteIPFragStream((*listIter)->id);
    }
    if (!m_IPFragStreamMap.empty()) { // 到这里数据出错了
        LOG_ERROR("IPFragStream leakage!!!");
    }
}

void IPReassembly::set_cfg(uint32_t ip_stream_size, uint32_t frag_cache_size, uint32_t timeout, bool detect_ooo)
{
    if (ip_stream_size == 0 || frag_cache_size == 0 || timeout == 0) { // 非法参数
        return;
    }
    m_cfg.max_stream_num = ip_stream_size;
    m_cfg.max_frag_cache_num = frag_cache_size;
    m_cfg.timeout = timeout;
    m_cfg.detect_ooo_enable = detect_ooo;
}

Packet* IPReassembly::process_packet(Packet* pkt, uint32_t &status, ProtocolType end_proto)
{
    aging_proc(); // 老化处理

    status = REASSEMBLY_ERROR; // 设置初值

    if (pkt == NULL || pkt->m_MetaData == NULL) {
        return NULL;
    }

    if (!pkt->isFragmentPacket()) {
        status = REASSEMBLY_NOT_FRAG;
        m_stat.not_frag_num++;
        return pkt;
    }
    m_stat.frag_rcv_num++;

    packet_meta_t *meta = pkt->m_MetaData;
    uint8_t *pkt_data = pkt->getRawPacketData();

    // 查找分片流是否存在
    IPFragStreamId id;
    getIPFragStreamId(meta, &id);
    IPFragStream *stream = getIPFragStream(id);
    if (stream == NULL) { // 新的ip分片流
        stream = addIPFragStream(id);
        if (stream == NULL) {
            return NULL;
        }
    }

    // 创建新的分片
    uint8_t *frag_data = pkt_data + meta->ip_offset + meta->ip_hdr_len;
    size_t frag_len = meta->ip_len - meta->ip_hdr_len;
    IPFragData *frag = new IPFragData(frag_data, frag_len, meta->ip_frag_offset, !(meta->ip_more_frag));
    if (frag == NULL || frag->data == NULL) {
        m_stat.frag_ceate_fail++;
        if (frag != NULL) {
            delete frag;
        }
        deleteIPFragStream(id);
        return NULL;
    }

    status = addIPFrag(stream, frag);
    if ( status != REASSEMBLY_CACHE_FRAG) {
        delete frag;
        if (status == REASSEMBLY_ERROR) {
            deleteIPFragStream(id);
        }
        return NULL;
    }

    // 所有分片报文都到达，进行报文重组,并重新解析
    if (is_complete(stream)) {
        status = REASSEMBLY_OK;
        Packet *new_pkt = reassemblyIPFragStream(stream, pkt, meta);
        // 重新解析报文,计算校验和
        new_pkt->parsePacket(end_proto);
        if (meta->ip_version == 4) {
            new_pkt->getLayerOfType<IPv4Layer>()->computeCalculateFields();
        }
        deleteIPFragStream(id);
        return new_pkt;
    }

    return NULL;
}

void IPReassembly::getIPFragStreamId(packet_meta_t *meta, IPFragStreamId *id)
{
    if (meta == NULL || id == NULL) {
        return;
    }

    id->source = meta->source;
    if (meta->is_vxlan) {
        id->vlanid = meta->vni;
    } else {
        id->vlanid = meta->vlanid;
    }
    id->sip = meta->sip;
    id->dip = meta->dip;
    id->fragid = meta->ip_id;
}

IPReassembly::IPFragStream* IPReassembly::getIPFragStream(IPFragStreamId &id)
{
    IPFragStream *stream = NULL;
    std::map<IPFragStreamId, IPStreamKeyValuePair>::iterator mapIter = m_IPFragStreamMap.find(id);
    
    if (mapIter == m_IPFragStreamMap.end()) { // 分片链不存在，创建分片链
        LOG_DEBUG("Got new packet with key[source:{} vlanid:{} sip:{} dip:{} pkt id:{}]",
            id.source, id.vlanid, id.sip.toString(), id.dip.toString(), id.fragid);
    }
    else { // 找到，刷新时间
        stream = mapIter->second.second;
        stream->update_time = getSystemRunTime();

        // 将节点移到老化链表头部
        if (mapIter->second.first != m_AgingList.begin()) {
            m_AgingList.erase(mapIter->second.first);
            mapIter->second.first = m_AgingList.insert(m_AgingList.begin(), stream);
        }
    }
    return stream;
}

IPReassembly::IPFragStream* IPReassembly::addIPFragStream(IPFragStreamId id)
{
    if (m_stat.stream_cur_num >= m_cfg.max_stream_num) {
        m_stat.stream_over_spec_num++;
        deleteOldestIPFragStream();
    }

    IPFragStream *stream = new IPFragStream(id);
    if (stream == NULL) {
        m_stat.stream_create_fail++;
        return NULL;
    }

    // 添加到老化链和map中
    auto list_iter = m_AgingList.insert(m_AgingList.begin(), stream); // 添加到链表头部
    IPStreamKeyValuePair pair(list_iter, stream);
    m_IPFragStreamMap.insert(std::make_pair(id, pair));
    m_stat.stream_cur_num++;
    m_stat.stream_total_num++;
    LOG_DEBUG("Add IP fragment stream[source:{} vlanid:{} sip:{} dip:{} pkt id:{}]",
            id.source, id.vlanid, id.sip.toString(), id.dip.toString(), id.fragid);
    return stream;
}

void IPReassembly::deleteIPFragStream(IPFragStreamId &id)
{
    auto mapIter = m_IPFragStreamMap.find(id);
    if (mapIter == m_IPFragStreamMap.end()) {
        LOG_ERROR("Delete IP fragment stream[source:{} vlanid:{} sip:{} dip:{} pkt id:{}] failed!", 
            id.source, id.vlanid, id.sip.toString(), id.dip.toString(), id.fragid);
        return;
    }

    IPFragStream *stream = mapIter->second.second;
    uint32_t fragCount = (uint32_t)stream->frag_list.size();
    auto listIter = mapIter->second.first;
    // 释放stream内存
    if (stream != NULL) {
        delete stream;
    }
    // 从map中删除
    m_IPFragStreamMap.erase(mapIter);

    // 从老化链中删除
    m_AgingList.erase(listIter);
    m_stat.stream_cur_num--;
    m_stat.frag_cache_num = m_stat.frag_cache_num - fragCount;

    LOG_DEBUG("Delete IP fragment stream[source:{} vlanid:{} sip:{} dip:{} pkt id:{}].", 
        id.source, id.vlanid, id.sip.toString(), id.dip.toString(), id.fragid);
}

void IPReassembly::print_ip_frag_streams()
{
    using namespace std;
    cout<<"IPFragStreams count:"<<endl;
    cout<<"list count:"<<m_AgingList.size()<<endl;
    cout<<"map count:"<<m_IPFragStreamMap.size()<<endl;
    map<IPFragStreamId, IPStreamKeyValuePair>::iterator iter = m_IPFragStreamMap.begin();
    uint32_t num = 0;
    while(iter != m_IPFragStreamMap.end()) {
        num++;
        IPFragStreamId id = iter->first;
        cout<<num<<":"<<"key("<<id.source<<","<<id.vlanid<<","<<id.fragid<<","
            <<id.sip.toString()<<","<<id.dip.toString()<<")"<<endl;
        iter++;
    }
}

void IPReassembly::deleteOldestIPFragStream()
{
    if(m_AgingList.empty()) {
        return;
    }
    // 删除链表尾部的分片流
    deleteIPFragStream(m_AgingList.back()->id);
}

uint32_t IPReassembly::addIPFrag(IPFragStream* stream, IPFragData *frag)
{
    if (m_stat.frag_cache_num >= m_cfg.max_frag_cache_num) {
        m_stat.frag_over_spec_num++;
        if (m_AgingList.back() == stream) {
            return REASSEMBLY_ERROR;
        }
        deleteOldestIPFragStream();
    }

    std::list<IPFragData*>::iterator pos;
    IPFragData *cur_node = NULL, *pre_node = NULL, *nextNode = NULL;
    
    // 遍历分片缓存链，找插入点
    for (pos = stream->frag_list.begin(); pos != stream->frag_list.end(); pos++) {
        cur_node = (*pos);
        if (cur_node->offset == frag->offset && cur_node->data_len == frag->data_len) { // 重复报文，丢弃
            m_stat.frag_dup_num++;
            return REASSEMBLY_DROP_FRAG;
        }

        if(cur_node->offset > frag->offset) { // 找到插入点
            nextNode = cur_node;
            break;
        }
        pre_node = cur_node;
    }

    // 重叠检查
    if (((pre_node != NULL) && (pre_node->offset + pre_node->data_len > frag->offset)) ||
        ((nextNode != NULL) && (frag->offset + frag->data_len > nextNode->offset))) {
        m_stat.frag_overlap_num++;
        return REASSEMBLY_ERROR;
    }

    // 乱序检测
    if (m_cfg.detect_ooo_enable) {
        (void)detectPktOutOfOrder(stream, frag, pos);
    }

	// 插入链表中
	stream->frag_list.insert(pos, frag);
	if (frag->offset == 0) {
		stream->rcv_first = true;
	} else if (frag->last) {
		stream->rcv_last = true;
		stream->total_size = frag->offset + frag->data_len;
	}
    stream->rcv_size += frag->data_len;
	m_stat.frag_cache_num++;
	return REASSEMBLY_CACHE_FRAG;
}

bool IPReassembly::detectPktOutOfOrder(IPFragStream* stream, IPFragData *frag,
    std::list<IPFragData*>::iterator pos)
{
    if (stream->expect_offset != frag->offset) { // 乱序报文
        m_stat.frag_ooo_num++;
        return true;
    } 

    // 刷新m_expectOffset
    stream->expect_offset = frag->offset + (uint16_t)frag->data_len;
    auto iter = pos;
    while(iter != stream->frag_list.end()) {
        if ((*iter)->offset == stream->expect_offset) {
            stream->expect_offset = (*iter)->offset + (uint16_t)(*iter)->data_len;
            iter++;
        } else {
            break;
        }
    }
    return false;
}

// 检查分片报文是否都已经到了
bool IPReassembly::is_complete(IPFragStream* stream)
{
    return (stream->rcv_first && stream->rcv_last && (stream->rcv_size == stream->total_size));
}

// 更新IP头部的报文长度和分片信息
void IPReassembly::updateIPHead(uint8_t version, uint8_t *ip_hdr, size_t hdr_len, size_t data_len, uint8_t proto)
{
    if (version == 4) {
        iphdr* ip4hdr = (iphdr*)(ip_hdr);
        ip4hdr->fragmentOffset = 0;
        ip4hdr->totalLength = htobe16(data_len + hdr_len);
    } else {
        ip6_hdr *ip6hdr = (ip6_hdr*)(ip_hdr);
        if (ip6hdr->nextHeader == IPv6Extension::IPv6Fragmentation) {
            ip6hdr->nextHeader = proto;
        } else {
            size_t offset = sizeof(ip6_hdr);
            while(offset < hdr_len) {
                ipv6_ext_base_header *extHdr = (ipv6_ext_base_header *)(ip_hdr + offset);
                if (extHdr->nextHeader == IPv6Extension::IPv6Fragmentation) {
                    extHdr->nextHeader = proto;
                    break;
                }
                offset += (extHdr->headerLen + 1) * 8;
            }
            if (offset >= hdr_len) {
                LOG_ERROR("IPReassembly IPv6 packet head error!");
            }
        }
        ip6hdr->payloadLength = htobe16(hdr_len + data_len - sizeof(ip6_hdr));
    }
}

// 重组报文
Packet* IPReassembly::reassemblyIPFragStream(IPFragStream* stream, Packet* pkt, packet_meta_t *meta)
{
    // 计算重组后的报文长度
    size_t hdr_len = meta->ip_offset + meta->ip_hdr_len;
    size_t pkt_len = stream->total_size + meta->ip_offset + meta->ip_hdr_len;
    if (meta->ip_version == 6) { // ipv6 需要减掉分片扩展头，固定长8字节,固定在IP头的最后面
        pkt_len -= 8;
        hdr_len -= 8;
    }

    //申请新报文的内存空间
    uint8_t *pkt_buf = new uint8_t[pkt_len];
    if (pkt_buf == NULL) {
        return NULL;
    }

    // 拷贝头部数据
    memcpy(pkt_buf, pkt->getRawPacketData(), hdr_len);
    // 刷新ip头部数据(报文长度，分片字段，ipv6需要刷新下一个协议的数据)
    updateIPHead(meta->ip_version, pkt_buf + meta->ip_offset, hdr_len - meta->ip_offset,
        stream->total_size, meta->ip_protocol);

    // 拷贝数据
    uint8_t *ip_data = pkt_buf + hdr_len;
    std::list<IPFragData*>::iterator pos = stream->frag_list.begin();
    IPFragData *fragNode = NULL;    
    for (; pos != stream->frag_list.end(); pos++) {
        fragNode = (*pos);
        memcpy(ip_data + fragNode->offset, fragNode->data, fragNode->data_len);
    }

    Packet *new_pkt = new Packet(pkt_buf, pkt_len, true, meta->link_type);
    new_pkt->m_MetaData = meta;
    new_pkt->clearMetaData();
    m_stat.stream_done_num++;
    return new_pkt;
}

void IPReassembly::aging_proc()
{
    uint64_t now = getSystemRunTime();
    if (now -  m_last_aging_time < 1000) { // 两次老化间隔需要大于1秒
        return;
    }
    m_last_aging_time = now;

    uint32_t count = 0;
    while (!m_AgingList.empty()) {        
        IPFragStream *stream = m_AgingList.back();
        if (stream->update_time + m_cfg.timeout > now) { // 老化时间还没到
            return;
        }

        // 删除老化的连接
        deleteIPFragStream(stream->id);
        m_stat.stream_timeout_num++;
        count++;
        if (count >= m_aging_count_once) {
            return;
        }
    };
}
}
